from __future__ import print_function

import csv
import numpy as np


def accumulate_study_results(ids, prob):
    sum_result = {}
    cnt_result = {}
    size = prob.shape[0]
    for i in range(size):
        study_id = ids[i]
        idx = int(study_id)
        if idx not in cnt_result:
            cnt_result[idx] = 0.
            sum_result[idx] = np.zeros((1, prob.shape[1]), dtype=np.float32)
        cnt_result[idx] += 1
        sum_result[idx] += prob[i, :]
    for i in cnt_result.keys():
        sum_result[i][:] /= cnt_result[i]
    sum_result = finalize(sum_result)
    return sum_result


def to_cdf(submission):
    for i in submission.keys():
        submission[i] = np.cumsum(submission[i])
    return submission


def correct_first_last(submission):
    for i in submission.keys():
        submission[i][0][0] = 0.
        submission[i][0][-1] = 1.
    return submission


def check_monotonic_increasing(submission):
    for i in submission.keys():
        for j in range(submission[i].shape[1] - 1):
            if submission[i][0][j] > submission[i][0][j + 1]:
                submission[i][0][j+1] = submission[i][0][j]
    return submission


def finalize(submission):
    submission = correct_first_last(submission)
    submission = check_monotonic_increasing(submission)
    return submission


def create_submission(systole_sub, diastole_sub, path):
    fi = csv.reader(open("dataset/csv/sample_submission_validate.csv"))
    f = open(path, "w")
    fo = csv.writer(f, lineterminator='\n')
    fo.writerow(fi.next())
    for line in fi:
        idx = line[0]
        key, target = idx.split('_')
        key = int(key)
        out = [idx]
        if key in systole_sub:
            if target == 'Diastole':
                out.extend(list(diastole_sub[key][0]))
            else:
                out.extend(list(systole_sub[key][0]))
        else:
            print('Miss {0}'.format(idx))
        fo.writerow(out)
    f.close()
